"""

Mobile robot motion planning sample with Dynamic Window Approach

author: Atsushi Sakai (@Atsushi_twi), Göktuğ Karakaşlı

"""

import math
from enum import Enum

import matplotlib.pyplot as plt
import numpy as np

show_animation = True


def dwa_control(x, config, goal, ob):
    """
    Dynamic Window Approach control
    """

    dw = calc_dynamic_window(x, config)

    u, trajectory = calc_control_and_trajectory(x, dw, config, goal, ob)

    return u, trajectory


class RobotType(Enum):
    circle = 0
    rectangle = 1


class Config:
    """
    simulation parameter class
    """

    def __init__(self):
        # robot parameter
        self.max_speed = 1  # [m/s]
        self.min_speed = -1  # [m/s]

        self.max_accel = 0.5  # [m/ss]

        self.v_resolution = 0.1  # [m/s]
        self.yaw_rate_resolution = self.v_resolution / 5.0  # [rad/s]
        self.dt = 0.1  # [s] Time tick for motion prediction
        self.predict_time = 3  # [s]
        self.to_goal_cost_gain = 1
        self.speed_cost_gain = 0.1
        self.obstacle_cost_gain = 1
        self.robot_type = RobotType.circle

        # if robot_type == RobotType.circle
        # Also used to check if goal is reached in both types
        self.robot_radius = 0.1  # [m] for collision check

        # if robot_type == RobotType.rectangle
        self.robot_width = 1.3  # [m] for collision check
        self.robot_length = 2 # [m] for collision check

    @property
    def robot_type(self):
        return self._robot_type

    @robot_type.setter
    def robot_type(self, value):
        if not isinstance(value, RobotType):
            raise TypeError("robot_type must be an instance of RobotType")
        self._robot_type = value


def motion(x, u, dt):
    """
    motion model
    u[0] -- 线速度
    u[1] -- 角速度

    x[2] -- 角度（rad)
    x[3] -- 速度（m/s）
    x[4] -- 角速度（rad/s）
    """

    x[2] += u[1] * dt
    x[0] += u[0] * math.cos(x[2]) * dt
    x[1] += u[0] * math.sin(x[2]) * dt
    x[3] = u[0]
    x[4] = u[1]

    return x


def calc_dynamic_window(x, config):
    """
    calculation dynamic window based on current state x
    """

    # Dynamic window from robot specification
    Vs = [config.min_speed, config.max_speed]

    # Dynamic window from motion model
    # 假设电机加速是瞬间的
    Vd = [x[3] - config.max_accel,
          x[3] + config.max_accel]

    #  [v_min, v_max, yaw_rate_min, yaw_rate_max]
    dw = [max(Vs[0], Vd[0]), min(Vs[1], Vd[1]),
          0, 0]

    r = 5.0
    dw[2] = min(-dw[1]/r, dw[1]/r, -dw[0]/r, dw[0]/r)
    dw[3] = max(-dw[1]/r, dw[1]/r, -dw[0]/r, dw[0]/r)

    return dw


def predict_trajectory(x_init, v, y, config):
    """
    predict trajectory with an input
    """

    x = np.array(x_init)
    trajectory = np.array(x)
    time = 0
    while time <= config.predict_time:
        x = motion(x, [v, y], config.dt)
        trajectory = np.vstack((trajectory, x))
        time += config.dt

    return trajectory


def calc_control_and_trajectory(x, dw, config, goal, ob):
    """
    calculation final input with dynamic window
    """

    x_init = x[:]
    min_cost = float("inf")
    best_u = [0.0, 0.0]
    best_trajectory = np.array([x])

    #防止np.arange取不到数据，np.arange(0.1, 0.1, 0.1)返回为空数组
    err = 0.0001

    # evaluate all trajectory with sampled input in dynamic window
    for v in np.arange(dw[0], dw[1]+err, config.v_resolution):
        for y in np.arange(dw[2], dw[3]+err, config.yaw_rate_resolution):

            if abs(v)<0.005:continue

            trajectory = predict_trajectory(x_init, v, y, config)

            # calc cost
            to_goal_cost = config.to_goal_cost_gain * calc_to_goal_cost(trajectory, goal)
            speed_cost = config.speed_cost_gain * (config.max_speed - trajectory[-1, 3])
            ob_cost = config.obstacle_cost_gain * calc_obstacle_cost(trajectory, ob, config)

            final_cost = to_goal_cost + speed_cost + ob_cost

            # search minimum trajectory
            if min_cost >= final_cost:
                min_cost = final_cost
                best_u = [v, y]
                best_trajectory = trajectory

    if abs(best_u[0]) < 0.01:
        print('oops')
    return best_u, best_trajectory


def calc_obstacle_cost(trajectory, ob, config):
    """
        calc obstacle cost inf: collision
    """
    ox = ob[:, 0] 
    oy = ob[:, 1]
    dx = trajectory[-1, 0] - ox #扩充到二维，这样做减法相当于路径中的每个点都和ox做减法
    dy = trajectory[-1, 1] - oy #只考虑终点
    r = np.hypot(dx, dy)

    def distance_to_line(x, y, A, B, C):
        return abs(A*x+B*y+C)/(math.sqrt(A**2+B**2))

    x0 = trajectory[-1, 0]
    y0 = trajectory[-1, 1]
    k0 = math.tan(trajectory[-1, 2]) if abs(trajectory[-1, 2])!=math.pi/2 else 1
    A0 = k0
    B0 = -1 if abs(trajectory[-1, 2])!=math.pi/2 else 0
    C0 = y0 - k0*x0

    k1 = -1/k0 if k0!=0 else 1
    A1 = k1
    B1 = -1 if k0!=0 else 0
    C1 = y0 - k1*x0

    dline0 = distance_to_line(ox, oy, A0, B0, C0) #平行方向的距离
    dline1 = distance_to_line(ox, oy, A1, B1, C1) #垂直方向的距离

    def in_body():
        ret = np.array(abs(dline0)<config.robot_width/2).any() and np.array(abs(dline1)<config.robot_length/2).any()
        return ret

    min_r = np.min(r)

    if in_body():
        return 100.0 / min_r
    else:
        return 1.0 / min_r  # OK


def calc_to_goal_cost(trajectory, goal):
    """
        calc to goal cost with angle difference
    """

    dx = goal[0] - trajectory[-1, 0]
    dy = goal[1] - trajectory[-1, 1]

    cost = abs(dy/dx)+math.hypot(dx,dy)

    return cost


def plot_arrow(x, y, yaw, length=0.5, width=0.1):  # pragma: no cover
    plt.arrow(x, y, length * math.cos(yaw), length * math.sin(yaw),
              head_length=width, head_width=width)
    plt.plot(x, y)


def plot_robot(x, y, yaw, config):  # pragma: no cover
    if config.robot_type == RobotType.rectangle:
        outline = np.array([[-config.robot_length / 2, config.robot_length / 2,
                             (config.robot_length / 2), -config.robot_length / 2,
                             -config.robot_length / 2],
                            [config.robot_width / 2, config.robot_width / 2,
                             - config.robot_width / 2, -config.robot_width / 2,
                             config.robot_width / 2]])
        Rot1 = np.array([[math.cos(yaw), math.sin(yaw)],
                         [-math.sin(yaw), math.cos(yaw)]])
        outline = (outline.T.dot(Rot1)).T
        outline[0, :] += x
        outline[1, :] += y
        plt.plot(np.array(outline[0, :]).flatten(),
                 np.array(outline[1, :]).flatten(), "-k")
    elif config.robot_type == RobotType.circle:
        circle = plt.Circle((x, y), config.robot_radius, color="b")
        plt.gcf().gca().add_artist(circle)
        out_x, out_y = (np.array([x, y]) +
                        np.array([np.cos(yaw), np.sin(yaw)]) * config.robot_radius)
        plt.plot([x, out_x], [y, out_y], "-k")


def main(gx=10, gy=14, ob=None, robot_type=RobotType.circle):
    print(__file__ + " start!!")
    # initial state [x(m), y(m), yaw(rad), v(m/s), omega(rad/s)]
    x = np.array([0.0, 0.0, 0.0, 0.0, 0.0])
    # goal position [x(m), y(m)]
    goal = np.array([gx, gy])
    # obstacles [x(m) y(m), ....]
    if ob is None:
        ob = np.array([[-1, -1],
                    [0, 2],
                    [4.0, 2.0],
                    [5.0, 4.0],
                    [5.0, 5.0],
                    [5.0, 6.0],
                    [5.0, 9.0],
                    [8.0, 9.0],
                    [7.0, 9.0],
                    [8.0, 10.0],
                    [9.0, 11.0],
                    [12.0, 13.0],
                    [12.0, 12.0],
                    [15.0, 15.0],
                    [13.0, 13.0]
                    ])

    # input [forward speed, yaw_rate]
    config = Config()
    config.robot_type = robot_type
    trajectory = np.array(x)

    tick = 0

    while True:

        tick+=1

        u, predicted_trajectory = dwa_control(x, config, goal, ob)
        x = motion(x, u, config.dt)  # simulate robot
        '''
        x[0] -- robot x坐标
        x[1] -- robot y坐标
        '''
        trajectory = np.vstack((trajectory, x))  # store state history

        print("[%d] v: %.5f, angular: %.5f, x: %.5f, y: %.5f"%(tick, u[0], u[1], x[0], x[1]))

        if show_animation:
            plt.cla()
            plt.axis("equal")
            # for stopping simulation with the esc key.
            plt.gcf().canvas.mpl_connect(
                'key_release_event',
                lambda event: [exit(0) if event.key == 'escape' else None])
            plt.plot(predicted_trajectory[:, 0], predicted_trajectory[:, 1], "-g")
            plt.plot(x[0], x[1], "xr")
            plt.plot(goal[0], goal[1], "xb")
            plt.plot(ob[:, 0], ob[:, 1], "ok", markersize=1)
            plot_robot(x[0], x[1], x[2], config)
            plot_arrow(x[0], x[1], x[2])
            plt.grid(True)
            plt.pause(0.0001)

        # check reaching goal
        dist_to_goal = math.hypot(x[0] - goal[0], x[1] - goal[1])
        if dist_to_goal <= 0.2:
            print("Goal!!")
            break

    print("Done")
    if show_animation:
        plt.plot(trajectory[:, 0], trajectory[:, 1], "-r")
        plt.pause(0.0001)

    plt.show()


if __name__ == '__main__':
    main(robot_type=RobotType.rectangle)
